use core::alloc::{GlobalAlloc, Layout};
use std::sync::Mutex;

// 使用静态内存池
const HEAP_SIZE: usize = 1024 * 1024; // 1MB 堆

// 静态内存池，作为我们的堆
static mut HEAP: [u8; HEAP_SIZE] = [0; HEAP_SIZE];

static CURRENT_OFFSET: Mutex<usize> = Mutex::new(0);

unsafe fn mi_malloc(size: usize, _alignment: usize) -> *mut u8 {
    let mut current = CURRENT_OFFSET.lock().unwrap();
    if *current == 0 {
        *current = HEAP.as_ptr() as usize;
    }
    if *current % 4 != 0 {
        *current += 4 - *current % 4;
    }
    let start = *current;

    let end = start + size;
    // 检查内存池是否有足够空间
    if end > HEAP.as_ptr() as usize + HEAP_SIZE {
        return std::ptr::null_mut(); // 返回空指针，表示分配失败
    }

    // 更新当前分配位置
    *current = end;
    start as *mut u8
}

unsafe fn mi_free(ptr: *mut u8, size: usize) {
    let mut current = CURRENT_OFFSET.lock().unwrap();
    *current = *current - size;
    if *current < ptr as usize {
        *current = ptr as usize;
        return;
    }
}

pub struct MiMalloc;

unsafe impl GlobalAlloc for MiMalloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        mi_malloc(layout.size(), layout.align()) as *mut u8
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        mi_free(ptr, layout.size());
    }
}

#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

fn main() {
    let x: Vec<i64> = vec![65535; 1000]; // 在堆上分配一个包含 1000 个元素的向量
    let x_ptr = std::ptr::addr_of!(x);
    println!("Allocated addr: {:p}", x_ptr);

    for i in 0..1000 {
        unsafe {
            println!("HEAP[{}] = {}", i, &HEAP[i]);
        }
    }
}
